# -*- coding: utf-8 -*-
# !/usr/bin/python3
"""
Author :      wu
Description :
"""

import tensorflow as tf
from tensorflow.keras import optimizers


# 几种不同方法求f(x) = a * x ** 2 + b * x + c的导数
def func_1():
    x = tf.Variable(0, dtype=tf.float32, name="x")
    a = tf.constant(1, dtype=tf.float32)
    b = tf.constant(-2, dtype=tf.float32)
    c = tf.constant(1, dtype=tf.float32)

    with tf.GradientTape() as tape:
        y = a * tf.pow(x, 2) + b * x + c

    dy_dx = tape.gradient(y, x)
    print("一阶导数 dy_dx = ", dy_dx)

    # 也可以对常量进行求导
    with tf.GradientTape() as tape:
        tape.watch([a, b, c])
        y = a * tf.pow(x, 2) + b * x + c
    dy_dx, dy_da, dy_db, dy_dc = tape.gradient(y, [x, a, b, c])
    print("常量求导：dy_dx = {}, dy_da = {}, dy_db = {}, dy_dc = {}".
          format(dy_dx, dy_da, dy_db, dy_dc))

    # 求二阶导数
    with tf.GradientTape() as tape2:
        with tf.GradientTape() as tape1:
            y = a * tf.pow(x, 2) + b * x + c
        dy_dx = tape1.gradient(y, x)
    d2y_d2x = tape2.gradient(dy_dx, x)

    print("二阶导数 d2y_d2x = ", d2y_d2x)


# autograph实现
@tf.function
def func_2(x):
    x = tf.cast(x, tf.float32)
    a = tf.constant(1, dtype=tf.float32)
    b = tf.constant(-2, dtype=tf.float32)
    c = tf.constant(1, dtype=tf.float32)

    with tf.GradientTape() as tape:
        tape.watch(x)
        y = a * tf.pow(x, 2) + b * x + c
    dy_dx = tape.gradient(y, x)

    return dy_dx, y


# 几种不同方法求f(x) = a * x ** 2 + b * x + c的最小值
# apply_gradients
def get_min():
    x = tf.Variable(0, name="x", dtype=tf.float32)
    a = tf.constant(1, dtype=tf.float32)
    b = tf.constant(-2, dtype=tf.float32)
    c = tf.constant(1, dtype=tf.float32)

    optimizer = optimizers.SGD(learning_rate=0.01)
    for _ in range(1000):
        with tf.GradientTape() as tape:
            y = a * tf.pow(x, 2) + b * x + c
        dy_dx = tape.gradient(y, x)
        optimizer.apply_gradients(grads_and_vars=[(dy_dx, x)])

    print("apply_gradients :y = {}, x= {}".format(y, x))


# minimize
def get_min_2():

    x = tf.Variable(0, name="x", dtype=tf.float32)

    def f():
        a = tf.constant(1, dtype=tf.float32)
        b = tf.constant(-2, dtype=tf.float32)
        c = tf.constant(1, dtype=tf.float32)
        y = a * tf.pow(x, 2) + b * x + c
        return y

    optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)
    for _ in range(1000):
        optimizer.minimize(f, [x])

    print("minimize: y = {}, x = {}".format(f(), x))


# apply_gradients + autograph
def get_min_3():

    x = tf.Variable(0, dtype=tf.float32, name="x")
    optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

    @tf.function
    def minimize_f():
        a = tf.constant(1, dtype=tf.float32)
        b = tf.constant(-2, dtype=tf.float32)
        c = tf.constant(1, dtype=tf.float32)

        for _ in tf.range(1000):
            with tf.GradientTape() as tape:
                y = a * tf.pow(x, 2) + b * x + c
            dy_dx = tape.gradient(y, x)
            optimizer.apply_gradients(grads_and_vars=[(dy_dx, x)])

        y = a * tf.pow(x, 2) + b * x + c
        return y

    # tf.print("apply_gradients + autograph: y = {}, x = {}".format(minimize_f(), x))
    print("apply_gradients + autograph: y = {:5}, x = {}".format(minimize_f(), x))


# minimize + autograph
def get_min_4():

    x = tf.Variable(0, dtype=tf.float32, name="x")
    optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

    @tf.function
    def f():
        a = tf.constant(1, dtype=tf.float32)
        b = tf.constant(-2, dtype=tf.float32)
        c = tf.constant(1, dtype=tf.float32)
        y = a * tf.pow(x, 2) + b * x + c
        return y

    @tf.function
    def train(epoch):
        for _ in tf.range(epoch):
            optimizer.minimize(f, [x])
        return f()

    print("minimize + autograph: y = {}, x = {}".format(train(1000), x))


def main():

    func_1()
    res = func_2(tf.constant(0))
    print("autograph: dy_dx = {:}, x = {}".format(res[0], res[1]))

    get_min()
    get_min_2()
    get_min_3()
    get_min_4()


if __name__ == "__main__":
    main()
